/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "slam_ros/playable_bag.h"

#include "absl/memory/memory.h"
#include "slam_ros/node_constants.h"
#include "glog/logging.h"
#include "tf2_msgs/msg/tf_message.hpp"
#include <rosbag2_cpp/reader.hpp>
#include <rosbag2_storage/storage_options.hpp>
#include <rclcpp/rclcpp.hpp>

namespace slam_ros {

inline std::ostream& operator<<(std::ostream& os, const rclcpp::Time& time) {
  os << time.nanoseconds();
  return os;
}

PlayableBag::PlayableBag(
    const std::string& bag_filename, const int bag_id,
    const rclcpp::Time start_time, const rclcpp::Time end_time,
    const rclcpp::Duration buffer_delay,
    FilteringEarlyMessageHandler filtering_early_message_handler)
    : reader_(absl::make_unique<rosbag2_cpp::Reader>()),
      finished_(false),
      bag_id_(bag_id),
      bag_filename_(bag_filename),
      message_counter_(0),
      buffer_delay_(buffer_delay),
      filtering_early_message_handler_(
          std::move(filtering_early_message_handler)),
      duration_in_seconds_(0.0) {
  
  rosbag2_storage::StorageOptions storage_options;
  storage_options.uri = bag_filename;
  storage_options.storage_id = "sqlite3";
  
  rosbag2_cpp::ConverterOptions converter_options;
  converter_options.input_serialization_format = "cdr";
  converter_options.output_serialization_format = "cdr";

  reader_->open(storage_options, converter_options);
  
  // Calculate duration
  auto metadata = reader_->get_metadata();
  const auto starting_time = rclcpp::Time(std::chrono::duration_cast<std::chrono::nanoseconds>(
      metadata.starting_time.time_since_epoch()).count());
  const auto ending_time = starting_time + 
                         rclcpp::Duration::from_nanoseconds(metadata.duration.count());
  const_cast<double&>(duration_in_seconds_) = (ending_time - starting_time).seconds();
  
  // Get all topics
  for (const auto& topic_info : metadata.topics_with_message_count) {
    topics_.insert(topic_info.topic_metadata.name);
  }
  
  AdvanceUntilMessageAvailable();
}

rclcpp::Time PlayableBag::PeekMessageTime() const {
  CHECK(IsMessageAvailable());
  return rclcpp::Time(buffered_messages_.front()->time_stamp);
}

std::tuple<rclcpp::Time, rclcpp::Time> PlayableBag::GetBeginEndTime() const {
  auto metadata = reader_->get_metadata();
  const auto starting_time = rclcpp::Time(std::chrono::duration_cast<std::chrono::nanoseconds>(
      metadata.starting_time.time_since_epoch()).count());
  const auto ending_time = starting_time + 
                          rclcpp::Duration::from_nanoseconds(metadata.duration.count());
  return std::make_tuple(starting_time, ending_time);
}

std::shared_ptr<rosbag2_storage::SerializedBagMessage> PlayableBag::GetNextMessage(
    slam_ros_msgs::msg::BagfileProgress* progress) {
  CHECK(IsMessageAvailable());
  auto msg = buffered_messages_.front();
  buffered_messages_.pop_front();
  AdvanceUntilMessageAvailable();
  
  auto msg_time = rclcpp::Time(msg->time_stamp);
  auto metadata = reader_->get_metadata();
  const auto starting_time = rclcpp::Time(std::chrono::duration_cast<std::chrono::nanoseconds>(
      metadata.starting_time.time_since_epoch()).count());
  double processed_seconds = (msg_time - starting_time).seconds();
  
  if ((message_counter_ % 10000) == 0) {
    LOG(INFO) << "Processed " << processed_seconds << " of "
              << duration_in_seconds_ << " seconds of bag " << bag_filename_;
  }

  if (progress) {
    progress->current_bagfile_name = bag_filename_;
    progress->current_bagfile_id = bag_id_;
    progress->total_messages = reader_->get_metadata().message_count;
    progress->processed_messages = message_counter_;
    progress->total_seconds = duration_in_seconds_;
    progress->processed_seconds = processed_seconds;
  }

  return msg;
}

bool PlayableBag::IsMessageAvailable() const {
  return !buffered_messages_.empty() &&
         (rclcpp::Time(buffered_messages_.front()->time_stamp) <
          rclcpp::Time(buffered_messages_.back()->time_stamp) - buffer_delay_);
}

int PlayableBag::bag_id() const { return bag_id_; }

void PlayableBag::AdvanceOneMessage() {
  CHECK(!finished_);
  if (!reader_->has_next()) {
    finished_ = true;
    return;
  }
  
  auto msg = reader_->read_next();
  if (!filtering_early_message_handler_ ||
      filtering_early_message_handler_(msg)) {
    buffered_messages_.push_back(msg);
  }
  ++message_counter_;
}

void PlayableBag::AdvanceUntilMessageAvailable() {
  if (finished_) {
    return;
  }
  do {
    AdvanceOneMessage();
  } while (!finished_ && !IsMessageAvailable());
}

PlayableBagMultiplexer::PlayableBagMultiplexer(const rclcpp::NodeOptions& options)
    : Node("playable_bag_multiplexer", options) {
  bag_progress_pub_ = this->create_publisher<slam_ros_msgs::msg::BagfileProgress>(
      "bagfile_progress", 10);
  this->declare_parameter<double>("bagfile_progress_pub_interval", 10.0);
  progress_pub_interval_ = this->get_parameter("bagfile_progress_pub_interval").as_double();
}

void PlayableBagMultiplexer::AddPlayableBag(PlayableBag playable_bag) {
  for (const auto& topic : playable_bag.topics()) {
    topics_.insert(topic);
  }
  playable_bags_.push_back(std::move(playable_bag));
  CHECK(playable_bags_.back().IsMessageAvailable());
  next_message_queue_.emplace(
      BagMessageItem{playable_bags_.back().PeekMessageTime(),
                     static_cast<int>(playable_bags_.size() - 1)});
  bag_progress_time_map_[playable_bags_.back().bag_id()] = this->now();
}

bool PlayableBagMultiplexer::IsMessageAvailable() const {
  return !next_message_queue_.empty();
}

std::tuple<std::shared_ptr<rosbag2_storage::SerializedBagMessage>, int, bool>
PlayableBagMultiplexer::GetNextMessage() {
  CHECK(IsMessageAvailable());
  const int current_bag_index = next_message_queue_.top().bag_index;
  PlayableBag& current_bag = playable_bags_.at(current_bag_index);
  slam_ros_msgs::msg::BagfileProgress progress;
  auto msg = current_bag.GetNextMessage(&progress);
  
  const bool publish_progress =
      current_bag.finished() ||
      (this->now() - bag_progress_time_map_[current_bag.bag_id()]).seconds() >=
          progress_pub_interval_;
  
  if (bag_progress_pub_->get_subscription_count() > 0 && publish_progress) {
    progress.total_bagfiles = playable_bags_.size();
    if (current_bag.finished()) {
      progress.processed_seconds = current_bag.duration_in_seconds();
    }
    bag_progress_pub_->publish(progress);
    bag_progress_time_map_[current_bag.bag_id()] = this->now();
  }
  
  CHECK_EQ(
    rclcpp::Time(msg->time_stamp).nanoseconds(),
    next_message_queue_.top().message_timestamp.nanoseconds()
  );
  
  next_message_queue_.pop();
  
  if (current_bag.IsMessageAvailable()) {
    next_message_queue_.emplace(
        BagMessageItem{current_bag.PeekMessageTime(), current_bag_index});
  }
  
  return std::make_tuple(std::move(msg), current_bag.bag_id(),
                         !current_bag.IsMessageAvailable());
}

rclcpp::Time PlayableBagMultiplexer::PeekMessageTime() const {
  CHECK(IsMessageAvailable());
  return next_message_queue_.top().message_timestamp;
}

}  // namespace slam_ros